"""
Dumbo Benchmark Fabric Tasks - Simplified Version
"""

from fabric import task
import os
import time
import subprocess
from typing import Dict, Any


@task
def local(ctx, nodes=4, duration=30, batch_size=100):
    """
    Run local Dumbo benchmark
    
    Args:
        nodes: Number of nodes (default: 4)
        duration: Test duration in seconds (default: 30)
        batch_size: Transaction batch size (default: 100)
    """
    print(f"\n{'='*60}")
    print("Starting local Dumbo benchmark")
    print(f"{'='*60}")
    
    try:
        # Change to dumbo root directory
        root_dir = os.path.dirname(os.path.abspath(__file__))
        dumbo_dir = os.path.join(root_dir, '..')
        
        print(f"Root directory: {dumbo_dir}")
        print(f"Nodes: {nodes}, Duration: {duration}s, Batch size: {batch_size}")
        
        # Run the local test script
        script_path = os.path.join(dumbo_dir, 'scripts', 'run_local_dumbo.sh')
        cmd = f"bash {script_path} {nodes} {duration}"
        
        print(f"Running command: {cmd}")
        result = subprocess.run(cmd, shell=True, cwd=dumbo_dir, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("✓ Local benchmark completed successfully")
            print(result.stdout)
        else:
            print("✗ Local benchmark failed")
            print("STDOUT:", result.stdout)
            print("STDERR:", result.stderr)
            return 1
            
    except Exception as e:
        print(f"✗ Local benchmark failed: {e}")
        return 1


@task
def status(ctx):
    """
    Show current benchmark status and configuration
    """
    print(f"\n{'='*60}")
    print("Dumbo Benchmark Status")
    print(f"{'='*60}")
    
    try:
        # Check if data directory exists
        root_dir = os.path.dirname(os.path.abspath(__file__))
        dumbo_dir = os.path.join(root_dir, '..')
        data_dir = os.path.join(dumbo_dir, "data")
        
        if os.path.exists(data_dir):
            runs = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
            print(f"Found {len(runs)} test runs in data directory")
            
            if runs:
                latest_run = sorted(runs)[-1]
                print(f"Latest run: {latest_run}")
                
                # Check results
                results_file = os.path.join(data_dir, latest_run, "results", "result.csv")
                if os.path.exists(results_file):
                    print("✓ Latest run has results available")
                else:
                    print("⚠ Latest run has no results")
        else:
            print("No test runs found")
        
        # Check configuration files
        config_files = ["settings.json", "aws_config.yaml", "alibaba_config.yaml"]
        for config_file in config_files:
            config_path = os.path.join(root_dir, config_file)
            if os.path.exists(config_path):
                print(f"✓ Configuration file exists: {config_file}")
            else:
                print(f"⚠ Configuration file missing: {config_file}")
        
        print("✓ Status check completed")
        
    except Exception as e:
        print(f"✗ Status check failed: {e}")
        return 1


@task
def clean(ctx):
    """
    Clean up local test data and temporary files
    """
    print(f"\n{'='*60}")
    print("Cleaning up local test data")
    print(f"{'='*60}")
    
    try:
        import shutil
        
        # Clean data directory
        root_dir = os.path.dirname(os.path.abspath(__file__))
        dumbo_dir = os.path.join(root_dir, '..')
        data_dir = os.path.join(dumbo_dir, "data")
        
        if os.path.exists(data_dir):
            shutil.rmtree(data_dir)
            print("✓ Cleaned data directory")
        
        # Clean temporary files
        temp_files = ["*.log", "*.tmp", "*.pid"]
        for pattern in temp_files:
            os.system(f"rm -f {pattern}")
        
        print("✓ Cleanup completed successfully")
        
    except Exception as e:
        print(f"✗ Cleanup failed: {e}")
        return 1


@task
def demo(ctx):
    """
    Run a quick demo of the Dumbo benchmark
    """
    print(f"\n{'='*60}")
    print("Dumbo Benchmark Demo")
    print(f"{'='*60}")
    
    try:
        # Run a quick local test
        print("Running quick demo test...")
        result = local(ctx, nodes=4, duration=10, batch_size=50)
        
        if result == 0:
            print("✓ Demo completed successfully")
            print("\nNext steps:")
            print("1. Check results: ls -la ../data/")
            print("2. Run longer test: fab local --nodes=7 --duration=60")
            print("3. View help: fab <command> --help")
        else:
            print("✗ Demo failed")
            
    except Exception as e:
        print(f"✗ Demo failed: {e}")
        return 1


@task
def help_custom(ctx):
    """
    Show detailed help for Dumbo benchmark
    """
    print(f"\n{'='*60}")
    print("Dumbo Benchmark Help")
    print(f"{'='*60}")
    
    print("\nAvailable commands:")
    print("  fab local          - Run local benchmark")
    print("  fab status         - Show benchmark status")
    print("  fab clean          - Clean up test data")
    print("  fab demo           - Run quick demo")
    print("  fab help_custom    - Show this help")
    
    print("\nLocal benchmark options:")
    print("  --nodes=N          - Number of nodes (default: 4)")
    print("  --duration=N       - Test duration in seconds (default: 30)")
    print("  --batch_size=N     - Transaction batch size (default: 100)")
    
    print("\nExamples:")
    print("  fab local --nodes=7 --duration=60")
    print("  fab local --nodes=4 --duration=30 --batch_size=200")
    print("  fab demo")
    print("  fab status")
    print("  fab clean")
    
    print("\nFor more information, see README.md")

